import torch
import torch.nn as nn
import torch.functional as F

class MyNet(nn.Module):
    def __init__(self, input_dim, hidden_dim_1, hidden_dim_2, hidden_dim_3, out_dim):
        super(MyNet, self).__init__()
        self.layer01 = nn.Linear(input_dim, hidden_dim_1)
        self.bn1 = nn.BatchNorm1d(hidden_dim_1)  # 批归一化操作，用来减缓梯度消失
        self.layer02 = nn.Linear(hidden_dim_1, hidden_dim_2)
        self.bn2 = nn.BatchNorm1d(hidden_dim_2)
        self.layer03 = nn.Linear(hidden_dim_2, hidden_dim_3)
        self.bn3 = nn.BatchNorm1d(hidden_dim_3)
        self.layer04 = nn.Linear(hidden_dim_3, out_dim)

    def forward(self, x):
        x = F.relu(self.layer01(x))
        x = self.bn1(x)
        x = F.relu(self.layer02(x))
        x = self.bn2(x)
        x = F.relu(self.layer03(x))
        x = self.bn3(x)
        x = torch.sigmoid(self.layer04(x))
        return x


